from flask import Flask, request, jsonify
from flask_cors import CORS
import os
import torch
from PIL import Image
from torchvision import transforms
from model import convnext_tiny as create_model
import mysql.connector
from werkzeug.utils import secure_filename

app = Flask(__name__)
CORS(app)

# 数据库配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'mf4B0zc4bFQGQQRcNUze',
    'database': 'ai_chat_db'
}

# AI模型配置
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
num_classes = 6
img_size = 400
model = create_model(num_classes=num_classes).to(device)
model_weight_path = "./weights/best_model_gpu_version2.pth"
model.load_state_dict(torch.load(model_weight_path, map_location=device))
model.eval()

data_transform = transforms.Compose([
    transforms.Resize(int(img_size * 1.14)),
    transforms.CenterCrop(img_size),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x[:3, ...]),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

UPLOAD_FOLDER = 'bronze_class_uploads'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

@app.route('/api4/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    if file:
        filename = secure_filename(file.filename)
        file_path = os.path.join(UPLOAD_FOLDER, filename)
        file.save(file_path)
        return jsonify({'message': 'File uploaded successfully', 'file_path': file_path}), 200

@app.route('/api4/predict', methods=['POST'])
def predict():
    data = request.json
    image_path = data.get('image_path')
    if not image_path:
        return jsonify({'error': 'No image path provided'}), 400

    img = Image.open(image_path)
    img = data_transform(img)
    img = torch.unsqueeze(img, dim=0)

    with torch.no_grad():
        output = torch.squeeze(model(img.to(device))).cpu()
        predict = torch.softmax(output, dim=0)
        predict_cla = torch.argmax(predict).item()
        probability = predict[predict_cla].item()

    # 将结果保存到数据库
    conn = mysql.connector.connect(**db_config)
    cursor = conn.cursor()
    sql = "INSERT INTO predictions (image_path, predicted_class, probability) VALUES (%s, %s, %s)"
    values = (image_path, str(predict_cla), probability)
    cursor.execute(sql, values)
    conn.commit()
    cursor.close()
    conn.close()

    return jsonify({
        'predicted_class': str(predict_cla),
        'probability': probability
    }), 200

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=1234, debug=True)